package com.example.webcrawler.controller;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;

import java.io.IOException;
import java.util.HashMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

// 添加WebSocket监控控制器
@RestController
@RequestMapping("/api/monitor")
public class MonitorController {

    private static final Logger logger = LoggerFactory.getLogger(MonitorController.class);

    // 存储活跃的WebSocket会话
    private final List<WebSocketSession> sessions = new CopyOnWriteArrayList<>();

    // 存储监控数据
    private final AtomicInteger activeSessions = new AtomicInteger(0);
    private final AtomicInteger totalRequests = new AtomicInteger(0);
    private final Map<String, Object> latestData = new ConcurrentHashMap<>();

    @GetMapping
    public ResponseEntity<Map<String, Object>> getMonitorStatus() {
        Map<String, Object> response = new HashMap<>();
        response.put("activeSessions", activeSessions.get());
        response.put("totalRequests", totalRequests.get());
        response.put("latestData", latestData);
        response.put("success", true);
        return ResponseEntity.ok(response);
    }

    @PostMapping("/data")
    public ResponseEntity<Map<String, Object>> updateMonitorData(@RequestBody Map<String, Object> data) {
        // 更新监控数据
        latestData.clear();
        latestData.putAll(data);

        // 如果有WebSocket会话，广播数据
        broadcastData(data);

        Map<String, Object> response = new HashMap<>();
        response.put("success", true);
        return ResponseEntity.ok(response);
    }

    // 获取WebSocket处理器
    @GetMapping("/websocket")
    public WebSocketHandler getWebSocketHandler() {
        return new MonitorWebSocketHandler();
    }

    // WebSocket处理器类
    public class MonitorWebSocketHandler extends TextWebSocketHandler {
        @Override
        public void afterConnectionEstablished(WebSocketSession session) throws Exception {
            sessions.add(session);
            activeSessions.incrementAndGet();
            logger.info("WebSocket连接建立，当前活跃会话数: {}", activeSessions.get());

            // 发送初始数据
            Map<String, Object> initData = new HashMap<>();
            initData.put("type", "init");
            initData.put("activeSessions", activeSessions.get());
            initData.put("totalRequests", totalRequests.get());
            session.sendMessage(new TextMessage(new com.fasterxml.jackson.databind.ObjectMapper().writeValueAsString(initData)));
        }

        @Override
        protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
            // 处理客户端消息
            logger.debug("收到WebSocket消息: {}", message.getPayload());
        }

        @Override
        public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
            sessions.remove(session);
            activeSessions.decrementAndGet();
            logger.info("WebSocket连接关闭，当前活跃会话数: {}", activeSessions.get());
        }

        @Override
        public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
            logger.error("WebSocket传输错误", exception);
            sessions.remove(session);
        }
    }

    // 广播数据给所有连接的客户端
    private void broadcastData(Map<String, Object> data) {
        String jsonData;
        try {
            jsonData = new com.fasterxml.jackson.databind.ObjectMapper().writeValueAsString(data);
        } catch (Exception e) {
            logger.error("序列化数据失败", e);
            return;
        }

        for (WebSocketSession session : sessions) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(jsonData));
                }
            } catch (IOException e) {
                logger.error("发送WebSocket消息失败", e);
            }
        }
    }

    // 更新请求数量
    public void incrementRequestCount() {
        totalRequests.incrementAndGet();
    }
}
